﻿#define IncrementalGenerator

using System;
using System.Collections.Generic;
using System.Collections.Immutable;
using System.Diagnostics;
using System.IO;
using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.CSharp.Syntax;
using Microsoft.CodeAnalysis.Text;
using UIOC.CodeAnalysis.Diagnostics;
using UIOC.CodeAnalysis.Extensions;
using UIOC.CodeAnalysis.Helpers;

#if IncrementalGenerator
#else
#endif

namespace UIOC.CodeAnalysis.Generators;

[Generator]
internal sealed class UIOCGenerator : Generator {

    public const string Leading = $"// <auto-generated> This file was automatically generated by {nameof(UIOCGenerator)}. </auto-generated>";

#if IncrementalGenerator
#else
    public override void Initialize(GeneratorInitializationContext context) => context.RegisterForSyntaxNotifications(() => new UIOCReceiver());
#endif

    protected override void Execute(
#if IncrementalGenerator
        IncrementalGeneratorInitializationContext context,
        IncrementalValueProvider<ImmutableArray<ClassDeclarationSyntax>> candidateClasses,
        IncrementalValueProvider<ImmutableArray<InterfaceDeclarationSyntax>> candidateInterfaces
#else
        GeneratorExecutionContext context,
        ImmutableArray<ClassDeclarationSyntax> candidateClasses,
        ImmutableArray<InterfaceDeclarationSyntax> candidateInterfaces
#endif
    ) {
        // 1. 生成实现类
        var generatedImplementations = Generate(context, candidateClasses, ImplementationGenerator.GenerateImplementation);

        // 2. 生成自动装配类
        var generatedAutowireds = Generate(context, candidateClasses, AutowiredGenerator.GenerateAutowired);

        // 3. 生成 Bean 类
        var generatedBeans = Generate(context, candidateClasses, BeanGenerator.GenerateBean);

        // 4. 生成 Bean 成员类 依赖 Bean 类 与 实现类
        var generatedBeanMembers = Generate(context, generatedBeans, BeanMembersGenerator.GenerateBeanMembers);

        // 5. 生成类型注册（依赖 Bean）
        var generatedTypeRegistrations = Generate(context, generatedBeans, TypeRegistrationGenerator.GenerateTypeRegistration);

        // 注册所有生成的源代码
        Register<
            ImmutableDictionary<INamedTypeSymbol, ImplementationGenerator.GeneratedInfo>,
            KeyValuePair<INamedTypeSymbol, ImplementationGenerator.GeneratedInfo>
        >(context, generatedImplementations, kv => kv.Value.GeneratedSource);

        Register<
            ImmutableArray<AutowiredGenerator.GeneratedInfo>,
            AutowiredGenerator.GeneratedInfo
        >(context, generatedAutowireds, info => info.GeneratedSource);

        Register<
            ImmutableDictionary<INamedTypeSymbol, BeanGenerator.GeneratedInfo>,
            KeyValuePair<INamedTypeSymbol, BeanGenerator.GeneratedInfo>
        >(context, generatedBeans, kv => kv.Value.GeneratedSource);

        Register<
            ImmutableArray<BeanMembersGenerator.GeneratedInfo>,
            BeanMembersGenerator.GeneratedInfo
        >(context, generatedBeanMembers, info => info.GeneratedSource, info => info.Diagnostics);

        Register<
            ImmutableArray<TypeRegistrationGenerator.GeneratedInfo>,
            TypeRegistrationGenerator.GeneratedInfo
        >(context, generatedTypeRegistrations, info => info.GeneratedSource);
    }
}

public class GeneratedSource(string hintName, string assemblyName, string namespaceName, SourceText sourceText, CompilationUnitSyntax compilationUnitSyntax, string generatorName) {

    public const string FileExtension = ".g.cs";

    public string HintName { get; } = hintName;

    public string AssemblyName { get; } = assemblyName;

    public string NamespaceName { get; } = namespaceName;

    public string GeneratorName { get; } = generatorName;

    public SourceText SourceText { get; } = sourceText;

    public CompilationUnitSyntax CompilationUnitSyntax { get; } = compilationUnitSyntax;

    public (string FullFilePath, string FilePath, string FileNameWithoutExtension) GetGeneratedSourceFilePath() {
        var fileNameWithoutExtension = $"{HintName}__{GeneratorName}";
        var fileName = $"{fileNameWithoutExtension}{FileExtension}";
        var fileDirectory = string.IsNullOrEmpty(NamespaceName) ? AssemblyName : $"{AssemblyName}/{NamespaceName}";
        if (SourceOutputHelper.CanWriteToProjectPath && SourceOutputHelper.OutputSourceGenFiles) {
            var generatedCodePath = SourceOutputHelper.GetGeneratedCodePath();
            var saveToDirectory = $"{generatedCodePath}/{fileDirectory}";
            Directory.CreateDirectory(saveToDirectory);
            return ($"{saveToDirectory}/{fileName}".Replace('\\', '/'), $"{fileDirectory}/{fileName}".Replace('\\', '/'), fileNameWithoutExtension.Replace('\\', '/'));
        }
        return ($"Temp/GeneratedCode/{fileDirectory}/{fileName}".Replace('\\', '/'), $"{fileDirectory}/{fileName}".Replace('\\', '/'), fileNameWithoutExtension.Replace('\\', '/'));
    }
}

#if IncrementalGenerator
public abstract class Generator : GeneratorBase {

    public override void Initialize(IncrementalGeneratorInitializationContext context) {
        var candidateInterfaces = context.SyntaxProvider
        .CreateSyntaxProvider(
            predicate: static (node, _) => node is InterfaceDeclarationSyntax ids && ids.AttributeLists.Count > 0,
            transform: static (ctx, _) => (InterfaceDeclarationSyntax)ctx.Node
        )
        .Collect();

        var candidateClasses = context.SyntaxProvider
            .CreateSyntaxProvider(
                predicate: static (node, _) => node is ClassDeclarationSyntax cds && cds.AttributeLists.Count > 0,
                transform: static (ctx, _) => (ClassDeclarationSyntax)ctx.Node
            )
            .Collect();

        var additionalFiles = context.AdditionalTextsProvider.Collect();
        context.RegisterSourceOutput(additionalFiles, (_, files) => {
            SourceOutputHelper.Setup(files);
        });
#else

public abstract class SyntaxReceiver : ISyntaxReceiver {

    public List<InterfaceDeclarationSyntax> CandidateInterfaces { get; } = [];

    public List<ClassDeclarationSyntax> CandidateClasses { get; } = [];

    public abstract void OnVisitSyntaxNode(SyntaxNode syntaxNode);
}

internal sealed class UIOCReceiver : SyntaxReceiver {

    public override void OnVisitSyntaxNode(SyntaxNode syntaxNode) {
        if (syntaxNode is InterfaceDeclarationSyntax interfaceDeclaration && interfaceDeclaration.AttributeLists.Count > 0) {
            CandidateInterfaces.Add(interfaceDeclaration);
        } else if (syntaxNode is ClassDeclarationSyntax classDeclaration && classDeclaration.AttributeLists.Count > 0) {
            CandidateClasses.Add(classDeclaration);
        }
    }
}

public abstract class Generator : GeneratorBase {

    public override void Execute(GeneratorExecutionContext context) {
        if (!SourceGenHelper.IsBuildTime || !SourceGenHelper.ShouldRun(context.Compilation, context.CancellationToken)) {
            return;
        }

        if (context.SyntaxReceiver is not SyntaxReceiver receiver) {
            return;
        }

        var candidateInterfaces = receiver.CandidateInterfaces.ToImmutableArray();
        var candidateClasses = receiver.CandidateClasses.ToImmutableArray();

        SourceOutputHelper.Setup(context.AdditionalFiles);
#endif

        var stopwatch = Stopwatch.StartNew();

        try {

            Execute(context, candidateClasses, candidateInterfaces);

            stopwatch.Stop();
            LogInfo($"TIME : {stopwatch.ElapsedMilliseconds} ms");

        } catch (OperationCanceledException) {
            throw;
        } catch (Exception exception) {
            var errorMessage = exception.ToUnityPrintableString();
            errorMessage = "This error indicates a bug in the UIOC source generators. " +
                "We'd appreciate a bug report (GitHub -> Issue...). Thanks! |--| " +
                $"Error message: '{errorMessage}'";
            var errorCode = GeneratorInternalDiagnostic.DiagnosticId;
            var errorTitle = GeneratorInternalDiagnostic.Title;
            SourceOutputHelper.LogInfo(DiagnosticSeverity.Error, $"Error: {errorCode}, {errorTitle}, {errorMessage}");
            var rule = new DiagnosticDescriptor(errorCode, errorTitle, errorMessage, "Source Generator", DiagnosticSeverity.Error, true);
#if IncrementalGenerator
            context.RegisterSourceOutput(context.CompilationProvider, (spc, _) => {
                spc.ReportDiagnostic(Diagnostic.Create(rule, Location.None));
            });
#else
            context.ReportDiagnostic(Diagnostic.Create(rule, Location.None));
#endif
        }
    }

    protected abstract void Execute(
#if IncrementalGenerator
        IncrementalGeneratorInitializationContext context,
        IncrementalValueProvider<ImmutableArray<ClassDeclarationSyntax>> candidateClasses,
        IncrementalValueProvider<ImmutableArray<InterfaceDeclarationSyntax>> candidateInterfaces
#else
        GeneratorExecutionContext context,
        ImmutableArray<ClassDeclarationSyntax> candidateClasses,
        ImmutableArray<InterfaceDeclarationSyntax> candidateInterfaces
#endif
    );

    private void LogInfo(string message) {
        SourceOutputHelper.LogInfo(DiagnosticSeverity.Info, $"[{nameof(Generator)}] {message}");
    }
}

public abstract class GeneratorBase :
#if IncrementalGenerator
    IIncrementalGenerator
#else
    ISourceGenerator
#endif
{
    private void LogInfo(string message) {
        SourceOutputHelper.LogInfo(DiagnosticSeverity.Info, $"[{nameof(GeneratorBase)}] {message}");
    }

#if IncrementalGenerator
    public void AddSource(SourceProductionContext context, GeneratedSource source) {
#else
    public void AddSource(GeneratorExecutionContext context, GeneratedSource source) {
#endif
        if (source != null) {
            var (fullFilePath, filePath, _) = source.GetGeneratedSourceFilePath();
            var sourceFilePath = filePath.Replace(GeneratedSource.FileExtension, string.Empty).Replace("/", "__").Replace('.', '_') + GeneratedSource.FileExtension;
            context.AddSource(sourceFilePath, source.SourceText);
            try {
                LogInfo($"[File({filePath})] ...\n{string.Join("\n", source.SourceText.Lines)}");
                SourceOutputHelper.OutputSourceToFile(fullFilePath, source.SourceText.ToString);
                //var generatorFullName = GetType().FullName;
                //var assembly = GetType().Assembly;
                //SourceOutputHelper.OutputSourceToFile($"{SourceOutputHelper.ProjectPath.Replace("\\", "/")}/{assembly.GetName().Name}/{generatorFullName}/{sourceFilePath}", source.SourceText.ToString);
            } catch (ArgumentException) {
            }
        }
    }

#if IncrementalGenerator
    public abstract void Initialize(IncrementalGeneratorInitializationContext context);

    public IncrementalValueProvider<T> Generate<T>(
        IncrementalGeneratorInitializationContext context,
        Func<Compilation, T> func
    ) {
        return context.CompilationProvider
            .Select((compilation, _) => {
                return func(compilation);
            });
    }

    public IncrementalValueProvider<T> Generate<T, T1>(
        IncrementalGeneratorInitializationContext context,
        IncrementalValueProvider<T1> combine1,
        Func<Compilation, T1, T> func
    ) {
        return context.CompilationProvider
            .Combine(combine1)
            .Select((tuple, _) => {
                var (compilation, t1) = tuple;
                return func(compilation, t1);
            });
    }

    public IncrementalValueProvider<T> Generate<T, T1, T2>(
        IncrementalGeneratorInitializationContext context,
        IncrementalValueProvider<T1> combine1,
        IncrementalValueProvider<T2> combine2,
        Func<Compilation, T1, T2, T> func
    ) {
        return context.CompilationProvider
            .Combine(combine1)
            .Combine(combine2)
            .Select((tuple, _) => {
                var ((compilation, t1), t2) = tuple;
                return func(compilation, t1, t2);
            });
    }

    public void Register<T, T1>(
        IncrementalGeneratorInitializationContext context,
        IncrementalValueProvider<T> generated,
        Func<T1, GeneratedSource> getSource,
        Func<T1, ImmutableArray<Diagnostic>> getDiagnostics = null
    ) where T : IEnumerable<T1> {
        context.RegisterSourceOutput(generated, (spc, sources) => {
            foreach (var item in sources) {
                var source = getSource(item);
                var diagnostics = getDiagnostics?.Invoke(item);
                if (diagnostics.HasValue) {
                    foreach (var diagnostic in diagnostics.Value) {
                        spc.ReportDiagnostic(diagnostic);
                    }
                }
                AddSource(spc, source);
            }
        });
    }
#else
    public abstract void Initialize(GeneratorInitializationContext context);
    public abstract void Execute(GeneratorExecutionContext context);

    public T Generate<T>(
        GeneratorExecutionContext context,
        Func<Compilation, T> func
    ) {
        return func(context.Compilation);
    }

    public T Generate<T, T1>(
        GeneratorExecutionContext context,
        T1 combine1,
        Func<Compilation, T1, T> func
    ) {
        return func(context.Compilation, combine1);
    }

    public T Generate<T, T1, T2>(
        GeneratorExecutionContext context,
        T1 combine1,
        T2 combine2,
        Func<Compilation, T1, T2, T> func
    ) {
        return func(context.Compilation, combine1, combine2);
    }
    public void Register<T, T1>(
        GeneratorExecutionContext context,
        T generatedInfos,
        Func<T1, GeneratedSource> getSource,
        Func<T1, ImmutableArray<Diagnostic>> getDiagnostics = null
    ) where T : IEnumerable<T1> {
        foreach (var info in generatedInfos) {
            var source = getSource(info);
            var diagnostics = getDiagnostics?.Invoke(info);
            if (diagnostics.HasValue) {
                foreach (var diagnostic in diagnostics.Value) {
                    context.ReportDiagnostic(diagnostic);
                }
            }
            AddSource(context, source);
        }
    }
#endif
}